#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Kafka环境设置助手

帮助用户设置Kafka环境，包括创建主题等
"""

import os
import argparse
import logging
import subprocess
import sys
import time
from kafka.admin import KafkaAdminClient, NewTopic
from kafka.errors import TopicAlreadyExistsError, NoBrokersAvailable
from kafka_config import get_kafka_config

# 配置日志
logging.basicConfig(
    format='%(asctime)s - %(levelname)s: %(message)s',
    level=logging.INFO
)
logger = logging.getLogger('kafka_setup')

def check_kafka_running(bootstrap_servers):
    """检查Kafka是否正在运行"""
    from kafka import KafkaProducer
    try:
        producer = KafkaProducer(bootstrap_servers=bootstrap_servers)
        producer.close()
        return True
    except NoBrokersAvailable:
        return False
    except Exception as e:
        logger.error(f"连接Kafka时发生错误: {str(e)}")
        return False

def create_topics(config=None):
    """创建Kafka主题"""
    if config is None:
        config = get_kafka_config()
    
    bootstrap_servers = config['bootstrap_servers']
    raw_data_topic = config['raw_data_topic']
    prediction_topic = config['prediction_topic']
    
    # 检查Kafka是否运行
    if not check_kafka_running(bootstrap_servers):
        logger.error(f"无法连接到Kafka服务器: {bootstrap_servers}")
        logger.error("请确保Kafka服务已启动")
        return False
    
    try:
        # 创建管理客户端
        admin_client = KafkaAdminClient(
            bootstrap_servers=bootstrap_servers
        )
        
        # 要创建的主题列表
        topic_list = [
            NewTopic(
                name=raw_data_topic,
                num_partitions=3,  # 使用3个分区以支持并行消费
                replication_factor=1  # 单机模式使用1个副本
            ),
            NewTopic(
                name=prediction_topic,
                num_partitions=3,
                replication_factor=1
            )
        ]
        
        # 创建主题
        admin_client.create_topics(new_topics=topic_list, validate_only=False)
        logger.info(f"成功创建Kafka主题: {raw_data_topic}, {prediction_topic}")
        
        admin_client.close()
        return True
        
    except TopicAlreadyExistsError:
        logger.warning("主题已存在，跳过创建")
        return True
    except Exception as e:
        logger.error(f"创建Kafka主题失败: {str(e)}")
        return False

def print_kafka_quick_start():
    """打印Kafka快速启动指南"""
    guide = """
Kafka 快速启动指南
=================

1. 下载 Kafka
   访问 https://kafka.apache.org/downloads 下载最新版本的Kafka

2. 解压 Kafka
   tar -xzf kafka_2.13-3.5.0.tgz
   cd kafka_2.13-3.5.0

3. 启动 Zookeeper (终端1)
   bin/zookeeper-server-start.sh config/zookeeper.properties

4. 启动 Kafka 服务器 (终端2)
   bin/kafka-server-start.sh config/server.properties

5. 创建主题 (可以使用本脚本自动创建)
   - 原始数据主题: air_quality_raw
   - 预测结果主题: air_quality_prediction

6. 监控主题 (可选，终端3)
   bin/kafka-console-consumer.sh --bootstrap-server localhost:9092 --topic air_quality_raw --from-beginning

恭喜！现在Kafka环境已准备就绪，可以运行空气质量预测系统了。
"""
    print(guide)

def main():
    parser = argparse.ArgumentParser(description='Kafka环境设置助手')
    parser.add_argument('--server', default='localhost:9092', help='Kafka服务器地址')
    parser.add_argument('--raw-topic', default='air_quality_raw', help='原始数据主题名称')
    parser.add_argument('--prediction-topic', default='air_quality_prediction', help='预测结果主题名称')
    
    args = parser.parse_args()
    
    # 使用命令行参数修改配置
    config = get_kafka_config()
    config['bootstrap_servers'] = args.server
    config['raw_data_topic'] = args.raw_topic
    config['prediction_topic'] = args.prediction_topic
    
    print("正在设置Kafka环境...")
    
    # 检查Kafka是否运行
    if check_kafka_running(config['bootstrap_servers']):
        logger.info(f"成功连接到Kafka服务器: {config['bootstrap_servers']}")
        
        # 创建主题
        if create_topics(config):
            logger.info("Kafka环境设置完成！")
        else:
            logger.error("Kafka环境设置失败！")
    else:
        logger.error(f"无法连接到Kafka服务器: {config['bootstrap_servers']}")
        logger.info("提供Kafka快速启动指南:")
        print_kafka_quick_start()

if __name__ == "__main__":
    main() 